package net.cyue.ort.llm.sampling;

import net.cyue.ort.llm.data.Token;

import java.util.ArrayList;
import java.util.List;

/**
 * 贪心采样策略
 * 始终选择概率最高的token
 */
public class GreedySamplingStrategy implements SamplingStrategy {

    @Override
    public Token selectNextToken(float[] logits, List<Long> context, SamplingConfig config) {
        List<Token> candidates = selectTopNTokens(logits, context, config, 1);
        if (candidates.isEmpty()) {
            return new Token(0L);
        }
        return candidates.get(0);
    }

    @Override
    public List<Token> selectTopNTokens(float[] logits, List<Long> context, SamplingConfig config, int topN) {
        if (logits == null || logits.length == 0) {
            return List.of();
        }
        SamplingConfig effectiveConfig = config != null ? config : SamplingConfig.defaultConfig();
        double[] adjusted = SamplingUtils.toDoubleArray(logits);
        SamplingUtils.applyRepetitionPenalty(adjusted, context, effectiveConfig.getRepetitionPenalty());
        SamplingUtils.applyTemperature(adjusted, Math.max(effectiveConfig.getTemperature(), 1e-6f));

        List<SamplingUtils.Candidate> candidates = SamplingUtils.buildCandidates(
            adjusted,
            effectiveConfig.getTopK(),
            effectiveConfig.getTopP()
        );

        int limit = Math.min(topN, candidates.size());
        List<Token> tokens = new ArrayList<>(limit);
        for (int i = 0; i < limit; i++) {
            SamplingUtils.Candidate candidate = candidates.get(i);
            tokens.add(new Token(candidate.tokenId(), candidate.logit()));
        }
        return tokens;
    }
}


